import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, activations, losses, optimizers, metrics


class ConvBn(keras.Model):

    def __init__(self, filters, ksize=(3, 3), strides=(1, 1), padding='same', **kwargs):
        super().__init__(**kwargs)
        self.conv = layers.Conv2D(filters, ksize, strides, padding, use_bias=False)
        self.bn = layers.BatchNormalization()

    def call(self, inputs, training=None, mask=None):
        x = self.conv(inputs, training=training)
        x = self.bn(x, training=training)
        return x


class ResnetBlock(keras.Model):

    def __init__(self, ch, is_shrink=False, **kwargs):
        super().__init__(**kwargs)
        self.is_shrink = is_shrink
        if is_shrink:
            strides = (2, 2)
            self.residual = ConvBn(ch, (1, 1), strides)
        else:
            strides = (1, 1)
        self.conv1 = ConvBn(ch, (3, 3), strides)
        self.conv2 = ConvBn(ch, (3, 3))

    def call(self, inputs, training=None, mask=None):
        # print('inputs', inputs.shape)
        x = self.conv1(inputs, training=training)
        x = layers.ReLU()(x, training=training)
        # print('x after conv1 relu', x.shape)
        x = self.conv2(x, training=training)
        # print('x after conv2', x.shape)

        if self.is_shrink:
            res = self.residual(inputs, training=training)
        else:
            res = inputs
        # print('res', res.shape)
        x += res
        x = layers.ReLU()(x, training=training)
        return x


class Resnet18(keras.Model):

    def __init__(self, n_cls, ch=64, **kwargs):
        super().__init__(**kwargs)
        self.conv1 = ConvBn(ch)
        blocks = []
        for id_unit in range(4):
            for id_block in range(2):
                if 0 == id_block and 0 != id_unit:
                    is_shrink = True
                    ch *= 2
                else:
                    is_shrink = False
                block = ResnetBlock(ch, is_shrink)
                blocks.append(block)
        self.blocks = blocks
        self.avg = layers.GlobalAvgPool2D()
        self.fc = layers.Dense(n_cls, activation=activations.softmax)

    def call(self, inputs, training=None, mask=None):
        x = self.conv1(inputs, training=training)
        x = layers.ReLU()(x, training=training)
        for i, block in enumerate(self.blocks):
            # print(i)
            x = block(x)
        x = self.avg(x, training=training)
        x = self.fc(x, training=training)
        return x


if '__main__' == __name__:

    model = Resnet18(10)
    model.build(input_shape=(None, 32, 32, 3))
    model.summary()

    x = tf.zeros((4, 32, 32, 3), dtype=tf.float32)
    pred = model(x)
    print('pred', pred.shape)
